from torchvision.datasets import CIFAR10
from PIL import Image


class PreDataset(CIFAR10):
    def __getitem__(self, item):
        img, target = self.data[item], self.targets[item]
        img = Image.fromarray(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        if self.transform is not None:
            imgL = self.transform(img)
            imgR = self.transform(img)
            return imgL, imgR, target
        else:
            return img, target


if __name__ == "__main__":
    from config import train_transform

    train_data = PreDataset(root='./dataset', train=True, transform=train_transform, download=True)
    print(train_data[0])
